#include <arch/cpu.h>
#include <arch/gdt.h>
#include <sys/sysdef.h>
#include <mm/valloc.h>
#include <mm/mmap.h>
#include <mm/vma.h>
#include <sched.h>
#include <process.h>
#include <fs.h>
#include <elf.h>
#include <log.h>

int map_elf_seg(Elf32_Phdr* ph, file_t* file) {
    u32 prot = 0;
    if (ph->p_flags & PF_R) {
        prot |= PROT_READ;
    }
    if (ph->p_flags & PF_W) {
        prot |= PROT_WRITE;
    }
    if (ph->p_flags & PF_X) {
        prot |= PROT_EXEC;
    }

    uptr vaddr = ph->p_vaddr;
    struct mmap_args args = {
        .file_offset = ph->p_offset,
        .len = (ph->p_memsz + PG_SIZE - 1) & ~(PG_SIZE - 1),
        .prot = prot,
        .flags = MMAP_FIXED,
        .mm = cur_proc->mm
    };
    int errno = do_memmap(NULL, PG_ALIGN(vaddr), file, &args);
    return errno;
}

bool check_elf_vaild(Elf32_Ehdr* h) {
    if (h->e_ident[EI_MAG0] != EI_MAG0 || h->e_ident[EI_MAG1] != EI_MAG1 || h->e_ident[EI_MAG2] != EI_MAG2 || h->e_ident[EI_MAG3] != EI_MAG3) {
        return false;
    }
    if (h->e_ident[EI_CLASS] != ELFCLASS32) {
        return false;
    }
    if (h->e_type != ET_EXEC) {
        return false;
    }
    if (h->e_machine != EM_386) {
        return false;
    }
    if (h->e_phoff == 0 || h->e_phentsize == 0 || h->e_phnum == 0) {
        return false;
    }
    return true;
}

int parse_elf(file_t* file, uptr* entry) {
    Elf32_Ehdr* h;
    Elf32_Phdr* ph;
    u8* buf = vzalloc(512);

    file->ops->read(file->inode, buf, 512, 0);
    h = (Elf32_Ehdr*)buf;

    if (check_elf_vaild(h)) {
        return ENOENT;
    }

    size_t header_size = h->e_phoff + h->e_phentsize * h->e_phnum;

    if (header_size > 512) {
        vfree(buf);
        buf = vzalloc(header_size);
        file->ops->read(file->inode, buf, header_size, 0);
    }
    *entry = h->e_entry;
    ph = (Elf32_Phdr*)(buf + h->e_phoff);
    for (int i = 0; i < h->e_phnum; i++) {
        if (ph[i].p_type != PT_LOAD) {
            continue;
        }
        if (ph[i].p_align != PG_SIZE) {
            continue;
        }
        map_elf_seg(&ph[i], file);
    }
    return 0;
}

size_t command_buffer_size(char** argv) {
    if (!argv) return 0;

    size_t len = 0, count = 1;
    while (*argv) {
        len += strlen(*argv) + 1;
        argv++;
        count++;
    }
    return len + count * sizeof(uptr);
}

size_t command_count(char** argv) {
    size_t len = 0;
    while (*argv) {
        len++;
        argv++;
    }
    return (len + 1) * sizeof(uptr);
}

uptr copy_commands(uptr stk_top, uptr* argv) {
    char* p;
    size_t len;

    while ((p = (char*)(*argv))) {
        len = strlen(p) + 1;

        stk_top -= len;
        memcpy((void*)stk_top, p, len);
        *argv = stk_top;

        argv++;
    }

    return stk_top;
}

int transfer_commands(file_t* file, char** argv, char** envp) {
    uptr ustk_top = cur_proc->mm->ustack_top;
    u32 argv_len = 0, envp_len = 0;
    uptr argvp = 0, envpp = 0;

    uptr buf_size = command_buffer_size(argv) + command_buffer_size(envp);
    uptr* buffer = vzalloc(buf_size);
    uptr buf = (uptr)buffer;
    buf += 511;

    if (envp) {
        envp_len = command_count(envp);
        buf -= envp_len;
        envpp = buf;

        memcpy((void*)buf, envp, envp_len);
        buf = copy_commands(buf, (uptr*)buf);
        envp = (char**)envpp;
    }

    if (argv) {
        argv_len = command_count(argv);
        buf -= argv_len;

        memcpy((void*)buf, argv, argv_len);
    }

    argv_len += sizeof(uptr);
    buf -= sizeof(uptr);
    *(char**)buf = file->dentry->name.val;

    argvp = buf;
    buf = copy_commands(buf, (uptr*)buf);
    argv = (char**)argvp;

    if (envpp) {
        ustk_top -= envp_len;
        envpp = ustk_top;

        memcpy((void*)ustk_top, envp, envp_len);
        ustk_top = copy_commands(ustk_top, (uptr*)ustk_top);
    }

    if (argvp) {
        ustk_top -= argv_len;
        argvp = ustk_top;

        memcpy((void*)ustk_top, argv, argv_len);
        ustk_top = copy_commands(ustk_top, (uptr*)ustk_top);
    }

    vfree(buffer);

    ustk_top -= sizeof(uptr) * 3;
    uptr* args = (uptr*)ustk_top;
    
    args[0] = (argv_len - 1) / sizeof(uptr);
    args[1] = argvp;
    args[2] = envpp;

    cur_proc->ustack_top = ustk_top;

    return 0;
}

__DEF_SYSCALL3(int, execve, const char*, path, char**, argv, char**, envp) {
    int errno = 0;
    dentry_t* dentry;
    file_t* file;

    if ((errno = fs_path_walk(path, cur_proc->cwd, &dentry, NULL, 0))) {
        goto out;
    }

    if ((errno = fs_open(dentry, &file))) {
        goto out;
    }

    assert(cur_proc->mm);
    assert(cur_proc->ustack_top);

    uptr entry = 0;
    if ((errno = parse_elf(file, &entry))) {
        goto out;
    }
    if ((errno = transfer_commands(file, argv, envp))) {
        goto out;
    }

    assign_return_context(0, cur_proc, cur_proc->context.esp);

    cur_proc->pgf_count = 0;

    cur_proc->context.ctx->ctx.eip = entry;
    cur_proc->context.ctx->ctx.esp0 = cur_proc->ustack_top;
    cur_proc->context.ctx->ctx.ss0 = SS_R3_DATA;

    asm volatile("pushl %0\n"
        "jmp switch_to\n" ::"r"(cur_proc)
        : "memory");

out:
    return errno;
}